package com.moneycollect.example.classloader.loader;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.util.List;
import java.util.zip.ZipEntry;
import java.util.zip.ZipInputStream;

/**
 * @author xuhuiqian
 * @since 2022/5/25 11:29
 */
public class MyClassLoader extends ClassLoader {
    private List<String> encryptedFileList;

    public MyClassLoader(String name, ClassLoader parent) {
        super(name, parent);
    }

    public MyClassLoader() {
        super();
    }

    public void setEncryptedFileList(List<String> encryptedFiles) {
        encryptedFileList = encryptedFiles;
    }

    public List<String> getEncryptedFileList() {
        return encryptedFileList;
    }

    @Override
    protected Class<?> findClass(String name) throws ClassNotFoundException {
        byte[] classData = loadClassData(name);
        if (classData == null) {
            throw new ClassNotFoundException();
        } else {
            return defineClass(name, classData, 0, classData.length);
        }
    }

    private byte[] loadClassData(String className) {
        if (this.encryptedFileList == null || this.encryptedFileList.size() < 1) {
            return null;
        }
        for (var file : this.encryptedFileList) {
            byte[] data = null;
            // if the file is java class file, required the file name format is: <fullClassName>.class
            if (file.endsWith(".class") && file.endsWith(className + ".class")) {
                data = loadClassDataFromClassFile(file);
            } else if (file.endsWith(".jar")) {
                data = loadClassDataFromJarFile(className, file);
            }
            if (data != null && data.length > 5) {
                return data;
            }
        }
        return null;
    }

    private byte[] loadClassDataFromClassFile(String fileName) {
        return LoaderUtils.decryptFile(fileName);
    }

    /**
     * to avoid saving unzip jar files to disk, use ZipInputStream
     */
    private byte[] loadClassDataFromJarFile(String className, String file) {
        try {
            byte[] encryptJarBytes = LoaderUtils.decryptFile(file);
            ZipInputStream zipStream = new ZipInputStream(new ByteArrayInputStream(encryptJarBytes));
            ZipEntry entry;
            while ((entry = zipStream.getNextEntry()) != null) {
                if (!entry.isDirectory() && entry.getName().endsWith(".class")) {
                    var entryName = entry.getName();
                    // if entryName includes /classes/, it indicates that it is a boot jar file,
                    // and must remove xxx/classes/ for the match to className
                    var classesDirPart = "/classes/";
                    if (entryName.contains(classesDirPart)) {
                        entryName = entryName.substring(entryName.indexOf(classesDirPart) + classesDirPart.length());
                    }
                    entryName = entryName.replace(".class", "").replace("/", ".");
                    if (className.equals(entryName)) {
                        ByteArrayOutputStream outStream = new ByteArrayOutputStream();
                        byte[] buffer = new byte[1024];
                        int length;
                        while ((length = zipStream.read(buffer)) != -1) {
                            outStream.write(buffer, 0, length);
                        }
                        zipStream.close();
                        return outStream.toByteArray();
                    }
                }
            }
            zipStream.close();
            return null;
        } catch (IOException e) {
            throw new RuntimeException("read class data from encrypted jar file error:" + e.getMessage(), e);
        }
    }
}
